package com.yupi.yuaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.yupi.yuaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;


import java.util.List;
import java.util.stream.Collector;
import java.util.stream.Collectors;

@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent{


    // 可用的工具
    private ToolCallback[] availableTools;

    // 工具调用回复变量
    private ChatResponse  toolCallChatResponse;

    // 工具调用guan
    private final ToolCallingManager toolCallingManager;

    // 禁用内置的工具调用机制 ， 自己维护上下文
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        // 禁用 SpringAI的内置工具调用 自己维护选项和消息上下文。
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true) // true
                .build();
    }

    @Override
    public boolean think() {
        if(getNextStepPrompt() != null && !getNextStepPrompt().isEmpty()){
            UserMessage userMessage = new UserMessage(getNextStepPrompt());
            getMessages().add(userMessage);
        }
        List<Message> messagesList = getMessages();
        Prompt prompt = new Prompt(messagesList, chatOptions);

        try {
            // 获取带工具选项的响应
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();
            this.toolCallChatResponse = chatResponse;
            AssistantMessage assistantMessage = toolCallChatResponse.getResult().getOutput();
            List<AssistantMessage.ToolCall> toolCalls = assistantMessage.getToolCalls();
            String text = assistantMessage.getText();
            String toolCallInfo = toolCalls.stream()
                    .map(toolCall -> {
                        return String.format("工具名称: %s , 参数： %s",
                                toolCall.name(),
                                toolCall.arguments());
                    }).collect(Collectors.joining("\n"));
            log.info("工具调用信息： ", toolCallInfo);
            // 判断是否有工具调用
            if(toolCalls.isEmpty()){
                // 没有工具调用就记录消息
                getMessages().add(assistantMessage);
                return false;
            }
            // 需要工具调用
            return true;
        } catch (Exception e) {
            log.error(getName() + "思考过程遇到了问题： ", e.getMessage());
            getMessages().add(new AssistantMessage("处理时遇到了错误：" + e.getMessage()));
            return false;
        }
    }

    @Override
    public String act() {
        // 需要调用工具
        if(!toolCallChatResponse.hasToolCalls()){
            return "没有工具调用";
        }
        // 调用工具
        Prompt prompt = new Prompt(getMessages(), chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        // 记录消息上下文
        setMessages(toolExecutionResult.conversationHistory());
        // 当前工具调用结果
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        String results = toolResponseMessage.getResponses().stream()
                .map(response -> {
                    return "工具 " + response.name() + "完成了他的任务！ 结果：" + response.responseData();
                }).collect(Collectors.joining("\n"));

        boolean terminatedToolCalled = toolResponseMessage.getResponses().stream()
                .anyMatch(response -> "doTerminate".equals(response.name()));
        if(terminatedToolCalled){
            setState(AgentState.FINISHED);
        }

        log.info("工具执行结果：{}", results);
        return results;
    }

    @Override
    public void clean() {
    }
}
